from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import time
import torch
from progress.bar import Bar
from models.data_parallel import DataParallel
from utils.utils import AverageMeter

import numpy as np

class ModleWithLoss(torch.nn.Module):
  def __init__(self, model, loss):
    super(ModleWithLoss, self).__init__()
    self.model = model
    self.loss = loss
  
  def forward(self, batch,global_step,tb_writer):
    outputs = self.model(batch['input'])
    loss, loss_stats = self.loss(outputs, batch,global_step,tb_writer)
    return outputs[-1], loss, loss_stats

class BaseTrainer(object):
  def __init__(
    self, opt, model, optimizer=None):
    self.opt = opt
    self.optimizer = optimizer
    self.loss_stats, self.loss = self._get_losses(opt)
    self.model_with_loss = ModleWithLoss(model, self.loss)

  def set_device(self, gpus, chunk_sizes, device):
    if len(gpus) > 1:
      self.model_with_loss = DataParallel(
        self.model_with_loss, device_ids=gpus, 
        chunk_sizes=chunk_sizes).to(device)
    else:
      self.model_with_loss = self.model_with_loss.to(device)
    
    for state in self.optimizer.state.values():
      for k, v in state.items():
        if isinstance(v, torch.Tensor):
          state[k] = v.to(device=device, non_blocking=True)

  def run_epoch(self, phase, epoch, data_loader,**kwargs):

    logger = kwargs.get('logger')
    model_with_loss = self.model_with_loss
    if phase == 'train':
      model_with_loss.train()
    else:
      if len(self.opt.gpus) > 1:
        model_with_loss = self.model_with_loss.module
      model_with_loss.eval()
      torch.cuda.empty_cache()

    opt = self.opt

    tb_writer = kwargs['tb_writer']
    results = {}
    data_time, batch_time = AverageMeter(), AverageMeter()
    avg_loss_stats = {l: AverageMeter() for l in self.loss_stats if not l.startswith('img_id') and not l.startswith('batch_')}
    num_iters = len(data_loader) if opt.num_iters < 0 else opt.num_iters
    bar = Bar('{}/{}'.format(opt.task, opt.exp_id), max=num_iters)
    end = time.time()
    print('num_iters:',num_iters)

   
    for iter_id, batch in enumerate(data_loader):
      if iter_id >= num_iters:
        break
      global_step= num_iters*(epoch-1)+iter_id
      data_time.update(time.time() - end)

      for k in batch:
        if k != 'meta':
          batch[k] = batch[k].to(device=opt.device, non_blocking=True)    
      output, loss, loss_stats = model_with_loss(batch,global_step,tb_writer)

      tb_time=time.time()
      if opt.equal_loss:
        if phase == 'train':
          loss = (loss[0]*opt.master_batch_size + loss[1]*(opt.batch_size - opt.master_batch_size))/opt.batch_size
        else:
          loss =loss.mean()
      else:        
        loss = loss.mean()
        
      if phase == 'train':
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
      
      batch_time.update(time.time() - end)
      end = time.time()

      Bar.suffix = '{phase}: [{0}][{1}/{2}]|Tot: {total:} |ETA: {eta:} '.format(
        epoch, iter_id, num_iters, phase=phase,
        total=bar.elapsed_td, eta=bar.eta_td)

      for l in avg_loss_stats:
        __stats= loss_stats
        if tb_writer is not None:
          tb_writer.add_scalar(l,__stats[l].mean().item(),global_step=global_step, walltime=tb_time)
        _n = batch['input'].size(0)
        avg_loss_stats[l].update(
          __stats[l].mean().item(), _n)
        Bar.suffix = Bar.suffix + '|{} {:.4f} '.format(l, avg_loss_stats[l].avg)

      if not opt.hide_data_time:
        Bar.suffix = Bar.suffix + '|Data {dt.val:.3f}s({dt.avg:.3f}s) ' \
          '|Net {bt.avg:.3f}s'.format(dt=data_time, bt=batch_time)
      if opt.print_iter > 0:
        if iter_id % opt.print_iter == 0:
          print('{}/{}| {}'.format(opt.task, opt.exp_id, Bar.suffix)) 
      else:
        bar.next()
      
      if opt.debug > 0 and opt.debug <5:
        self.debug(batch, output, iter_id)
      
      if opt.test:
        self.save_result(output, batch, results)
      del output, loss, loss_stats
    
    if tb_writer is not None:
      for l in avg_loss_stats:
        tb_writer.add_scalar('epoch_avg/{}'.format(l),avg_loss_stats[l].avg,global_step=global_step,walltime=time.time())
    bar.finish()
    ret = {k: v.avg for k, v in avg_loss_stats.items()}
    ret['time'] = bar.elapsed_td.total_seconds() / 60.
    return ret, results
  
  def debug(self, batch, output, iter_id):
    raise NotImplementedError

  def save_result(self, output, batch, results):
    raise NotImplementedError

  def _get_losses(self, opt):
    raise NotImplementedError
  
  def val(self, epoch, data_loader,**kwargs):
    return self.run_epoch('val', epoch, data_loader,**kwargs)

  def train(self, epoch, data_loader,**kwargs):
    return self.run_epoch('train', epoch, data_loader,**kwargs)

  def wta_stat(self,epoch, data_loader):

    stats = []
    model_with_loss = self.model_with_loss
    if len(self.opt.gpus) > 1:
      model_with_loss = self.model_with_loss.module
    model_with_loss.eval()
    torch.cuda.empty_cache()

    opt = self.opt
    results = {}
    data_time, batch_time = AverageMeter(), AverageMeter()
    avg_loss_stats = {l: AverageMeter() for l in self.loss_stats}
    num_iters = len(data_loader) if opt.num_iters < 0 else opt.num_iters
    end = time.time()
    for iter_id, batch in enumerate(data_loader):
      if iter_id >= num_iters:
        break
      global_step= num_iters*(epoch-1)+iter_id
      data_time.update(time.time() - end)

      for k in batch:
        if k != 'meta':
          batch[k] = batch[k].to(device=opt.device, non_blocking=True)    
      _,_,loss_stats = model_with_loss(batch,global_step,None)
      batch_time.update(time.time() - end)
      end = time.time()

    return stats